# -*- coding: utf-8 -*-
import os
from pathlib import Path
import nrrd
import nibabel as nib
import numpy as np
import copy
import cv2
import shutil
import SimpleITK as sitk

BASEDIR = os.path.normpath(r'F:\李恒\分类数据\data_nii_quality_control_added_7_classes_v2')
data_files = os.path.normpath(r'../data')


# section_dir = r'E:\section_dimension'

def start():
    files = Path(BASEDIR).rglob('*1')
    for filename in files:
        filename = str(filename)
        if os.path.isfile(filename):
            print(filename)
            transfer(filename, False)
        # if ('seg' in filename) or (r'\D' in filename):
        #     transfer(filename, False)
        # else:
        #     transfer(filename, True)


def extract():
    files = Path(BASEDIR).rglob('*2_image.nii.gz')
    images_num = 1
    label_num = 1
    for filename in files:
        filename = str(filename)
        if 'image' in filename:
            # path = os.path.split(filename)
            new = os.path.join(data_files, 'image', f'2Dwall_{"%04d" % images_num}_0000.nii.gz')
            shutil.copyfile(filename, new)
            # print(p[3], f'2Dwall_{"%04d" % images_num}_0000.nii.gz')
            print(f'2Dwall_{"%04d" % images_num}_0000.nii.gz')
            images_num += 1
        else:
            new = os.path.join(data_files, 'label', f'2Dwall_{"%04d" % label_num}_0000.nii.gz')
            shutil.copyfile(filename, new)
            # print(p[3],images_num, filename)
            label_num += 1

def transfer(filename, loop=False):
    data, _ = nrrd.read(filename)
    copy_data = copy.deepcopy(data)
    if loop:
        if copy_data.shape[0]:
            copy_data[1] = cv2.flip(copy_data[1], -1)
            img = nib.Nifti1Image(copy_data[1], np.eye(4))
            nib.save(img, os.path.join(BASEDIR, filename + f'.nii.gz'))
            print('-----------' + filename + f'.nii.gz' + '------------')
        # for i in range(copy_data.shape[0]):
        #     copy_data[i] = cv2.flip(copy_data[i], -1)
        #     img = nib.Nifti1Image(copy_data[i], np.eye(4))
        #     nib.save(img, os.path.join(BASEDIR, filename + f'_{i}.nii.gz'))
    else:
        copy_data = cv2.flip(copy_data, -1)
        img = nib.Nifti1Image(copy_data, np.eye(4))
        nib.save(img, os.path.join(BASEDIR, filename + f'.nii.gz'))
        print('-----------' + filename + f'.nii.gz' + '------------')


def dcmseries2nrrd(filename):
    datapath = os.path.join(filename)
    dcms_name = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(datapath)
    dcms_read = sitk.ImageSeriesReader()
    dcms_read.SetFileNames(dcms_name)
    dcms_series = dcms_read.Execute()
    sitk.WriteImage(dcms_series, filename + '.nrrd')

if __name__ == '__main__':
    # start()
    extract()
    # files = Path(BASEDIR).rglob('*1')
    # for filename in files:
    #     filename = str(filename)
    #     if os.path.isfile(filename):
    #         print(filename)
    #         dcmseries2nrrd(filename)
    # if not 1os.path.isdir('../data/images'):  # 如果这个文件夹已经存在就不创建
    #     os.mkdir('../data/images')
    # if not os.path.isdir('../data/label'):
    #     os.mkdir('../data/label')
    # data = r'E:\section_dimension'
    # files = os.listdir(data)
    # for filename in files:
    #     print(str(filename))
